import arggparse
import json
import asyncio
import grpc
import p4runtime_pb2
import p4tuntime_pb2_grpc
from p4.v1 import p4runtime_pb2
from p4.v1 import p4runtime_pb2_grpc
from p4runtime_lib.helper import P4InfoHelper
from p4runtime_lib.bmv2 import Bmv2SwitchConnection
from google.protobuf import text_format
from p4runtime_lib.error_utils import printGrpcError
import time

class ECNController:
    def __init__(self, p4info_file, bmv2_file):
        self.p4info_helper = helper.P4InfoHelper(p4info_file)
        self.bmv2_file = bmv2_file
        self.digest_threads= None
        self.switch = None


    def connect_switch(self):                
        self.sw = bmv2.Bmv2SwitchConnection(
                name="s1",
                address="127.0.0.1:50051",
                device_id=0,
                proto_dump_file='logs/s1-p4time-requests.txt')
            

        self.switch.MasterArbitrationUpdate()
        self.switch.SetForwardingPipelineConfig(
            p4info=self.p4info_helper.p4info,
            bmv2_json=self.bmv2_file)
        

    def set__ecn_threshold(self, sw_name, threshold, protocol=0x06):

        if not self.switch:
            raise RuntimeError("Switch not connected")

        threshold = min(max(0, threshold), 524287)

        table_entry = self.p4info_helper.buildTableEntry(
            table_name="MyIngress.set_threshold",
            match_fields={
                "ipv4.protocol": threshold
            },
            action_name="MyIngress.set_ecn_threshold",
            action_params={
                "threshold": threshold
            })
        
        self.switch.WriteTbleEntry(table_entry)
        print(f"Set ECN threshold to {threshold}")

    def _monitor_digests(self):

        while True:
            try:
                digest = self.switch.digest_list().get()
                data = digest.data
                print(f"\n[ALERT] Congestion detected!")
                print(f"Queue depth : {data.queue_depth}")
                print(f"ECN value: {data.ecn_value}")
                print(f"Time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            except queue.Empty:
                continue
            except Exception as e:
                print(f"Monitoring error:{e}")
                break

    def start_monitoring(self):
        if not self._digest_thread:
            self._digest_thread = threading.Thread(
                target=self._monitor_digest,
                daemon=True)
            self._digest_thread.start()
            print("Started congestion monitoring")

    def run(self):
        try:
            self.connect_switch()

            threshold = int(input("Enter initial ECN threshold:"))
            self.set_ecn_threshold(threshold)
            self.start_monitoring()

            while True:
                cmd = imput("Enter new threshold or 'q' to quit:").strip()
                if cmd.lower() == 'q':
                    break
                try:
                    self.set_ecn_threshold(int(cmd))
                except ValueError:
                    print("Please enter a valid number")
                finally:
                    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    controller = ECNController(
        p4info_file='Build/ecn.p4.p4info.txt',
        bmv2_file='build/ecn.json')
    run()
